from PySide6.QtCore import QObject, Signal
from loguru import logger
from concurrent.futures import ThreadPoolExecutor

class SlowProcessWorker(QObject):
    """
    A worker that processes a dictionary of slow functions and their callbacks in parallel.
    Each function is executed in a separate thread, and upon completion, a signal is emitted 
    with its result and the corresponding callback.
    """
    
    # Signal to emit when a task is done. 
    # It passes the result (any type) and the callback function.
    task_done = Signal(object, object)
    
    # Signal to emit when progress is updated
    progress_updated = Signal(int)

    def _handle_task_done(self, result, callback):
        """处理任务完成信号"""
        if callback:
            callback(result)

    def __init__(self, slow_tasks: dict):
        """
        Initializes the worker.
        
        Args:
            slow_tasks (dict): A dictionary where keys are the slow functions to execute,
                               and values are the callback functions to be called with the results.
        """
        super().__init__()
        self.slow_tasks = slow_tasks
        self.executor = ThreadPoolExecutor()
        self._running = False
        self._completed_tasks = 0
        self._total_tasks = len(slow_tasks)

        self.task_done.connect(self._handle_task_done)

    def update_progress(self, value):
        """更新进度 - 线程安全的方法"""
        # 发射进度更新信号到主线程
        self.progress_updated.emit(value)

    def start(self):
        """启动并行任务执行"""
        if self._running:
            return
            
        self._running = True
        
        try:
            for func, callback in self.slow_tasks.items():
                future = self.executor.submit(func)
                # 使用 lambda 捕获 callback，避免闭包问题
                future.add_done_callback(lambda f, cb=callback: self._on_task_done(f, cb))
        except Exception as e:
            logger.exception(f"启动任务时出错: {e}")
            self._running = False

    def _on_task_done(self, future, callback):
        """任务完成回调"""
        try:
            result = future.result()
            # 发送信号到主线程
            self.task_done.emit(result, callback)
        except Exception as e:
            logger.exception(f"A task in SlowProcessWorker resulted in an exception: {e}")
        finally:
            # 检查是否所有任务都完成了
            self._check_all_tasks_done()
    
    def _check_all_tasks_done(self):
        """检查所有任务是否完成，如果是则清理资源"""
        self._completed_tasks += 1
        if self._completed_tasks >= self._total_tasks:
            self._running = False
            logger.debug(f"All {self._total_tasks} tasks completed, worker stopped")
    
    def isRunning(self):
        """返回当前是否正在运行"""
        return self._running
    
    def shutdown(self):
        """关闭线程池并清理资源"""
        try:
            self.executor.shutdown(wait=True)
        except Exception as e:
            logger.warning(f"关闭线程池时出错: {e}")
        
        try:
            # 清理数据库连接
            from core.database import db_manager
            db_manager.close()
        except Exception as e:
            logger.warning(f"清理数据库连接时出错: {e}")
        finally:
            self._running = False